#!/usr/bin/env python
""" Group Project on Word Alignment

David Coles, Guan Gui, Kai Chen; 2010
$LastChangedBy$
"""

__version__ = "$Rev$"

from collections import defaultdict
try:
    import pylab
    __pylab_loaded__ = True
except ImportError:
    __pylab_loaded__ = False

class AlignedSent(object):
    """Class to store an alignment between two parallel sentences.
    """
    # Default encoding used by this class
    # N.B. This is a static class variable
    encoding = 'utf-8'

    def __init__(self, sent1, sent2, alignments=None):
        """Takes two lists of words, sent1 and sent2 which represent two 
        parallel sentences an optional iterateable object of tupples where 
        each tupple are 0-originated indexes of mappings between words.
        """
        self.sent1 = sent1
        self.sent2 = sent2
        self.sent1_index = None
        self.sent2_index = None

        if alignments:
            self.load_alignments(alignments)

    def load_alignments(self, alignments):
        """Takes a list of tupples of the form (pos_sent1, pos_sent2) that
        describe how the two sentences align up.
        """
        # my innovation, _ as dummy variable, seems can stop compiler
        # complaining about unused variables ;-)
        self.sent1_index = [set() for _ in self.sent1]
        self.sent2_index = [set() for _ in self.sent2]

        # Iterate through the aliments and create an index for each
        for (a_1, a_2) in alignments:
            try:
                self.sent1_index[a_1].add(a_2)
            except IndexError:
                raise IndexError(
                        "%d is not a valid alignment index for sentence 1"%
                        a_1)
            try:
                self.sent2_index[a_2].add(a_1)
            except IndexError:
                raise IndexError(
                        "%d is not a valid alignment index for sentence 2"%
                        a_2)

    def get_sent1_nth0_aligned_words(self, nth0):
        '''Return a list of words, in orignal order, in sentence 2 aligned with
        the nth (starts with 0) word in sentence 1
        '''
        return [self.sent2[i] for i in sorted(self.sent1_index[nth0])]

    def get_sent1_nth1_aligned_words(self, nth1):
        '''Return a list of words, in orignal order, in sentence 2 aligned with
        the nth (starts with 1) word in sentence 1
        '''
        return self.get_sent1_nth0_aligned_words(nth1 - 1)

    def get_sent2_nth0_aligned_words(self, nth0):
        '''Return a list of words, in orignal order, in sentence 1 aligned with
        the nth (starts with 0) word in sentence 2
        '''
        return [self.sent1[i] for i in sorted(self.sent2_index[nth0])]

    def get_sent2_nth1_aligned_words(self, nth1):
        '''Return a list of words, in orignal order, in sentence 1 aligned with
        the nth (starts with 1) word in sentence 2
        '''
        return self.get_sent2_nth0_aligned_words(nth1 - 1)

    def sent1_sets(self):
        """Returns a list of (word, set(mots)) tupples where word is a word in 
        sentence 1 and mots are the words it's aligned with in sentence 2.
        """
        word_sets = [set([self.sent2[i] for i in aset]) for aset in 
        self.sent1_index]
        return zip(self.sent1, word_sets)

    def sent2_sets(self):
        """Returns a list of (mot, set(words)) tupples where mot is a word in 
        sentence 2 and words are the words it's aligned with in sentence 1.
        """
        word_sets = [set([self.sent1[i] for i in aset]) for aset in 
                self.sent2_index]
        return zip(self.sent2, word_sets)

    def __repr__(self):
        """Returns a sentence alignment formatted in GIZA format.
        Line 1 is the first of the aligned sentences
        Line 2 is the second of the aligned sentence
        Line 3 is a space separated list of alignment indexes 'idx1-idx2'
        """
        lines = []
        # First 2 lines are just each sentences words separated by whitespace
        lines.append(" ".join(self.sent1))
        lines.append(" ".join(self.sent2))

        # Construct the alignment pairs by iterating through one of the 
        # alignment indexes.
        alignments = []
        for idx1 in range(len(self.sent1_index)):
            for idx2 in self.sent1_index[idx1]:
                alignments.append("%d-%d"%(idx1, idx2))
        lines.append(" ".join(alignments))

        # Return the 3 GIZA lines
        return "\n".join(lines).encode(AlignedSent.encoding)

    def gui_repr(self):
        """Generate a GUI to represent the sentence alignments
        """
        if __pylab_loaded__:
            fig_width = max(len(self.sent1), len(self.sent2)) + 1
            fig_height = 3
            pylab.figure(figsize=(fig_width*0.8, fig_height*0.8), facecolor='w')
            pylab.box(on=False)
            pylab.subplots_adjust(left=0, right=1, bottom=0, top=1)
            pylab.xlim(-1, fig_width - 1)
            pylab.ylim(0, fig_height)
            pylab.xticks([])
            pylab.yticks([])
            for i in xrange(len(self.sent1)):
                pylab.text(i, 2.5, self.sent1[i], ha = 'center', va = 'center',
                        rotation=30)
                if len(self.sent1_index[i]) > 0:
                    pylab.arrow(i, 2.5, 0, -0.5, color='r', alpha=0.3, lw=2)
                for j in self.sent1_index[i]:
                    pylab.arrow(i, 2, j - i, -1, color='r')
            for i in xrange(len(self.sent2)):
                pylab.text(i, 0.5, self.sent2[i], ha = 'center', va = 'center',
                        rotation=30)
                if len(self.sent2_index[i]) > 0:
                    pylab.arrow(i, 0.5, 0, 0.5, color='r', alpha=0.3, lw=2)
            pylab.draw()

    @staticmethod
    def block_and_view_gui_repr():
        """Block the program to view previous spawned Repr GUIs
        """
        if __pylab_loaded__:
            pylab.show()

def gen_dict(aligned_sents):
    """Generate a bi-lingual dictionary from a list of AlignedSents.
    the return value is in format of (en_to_fr_dict, fr_to_en_dict);
    the dictionary is potentially big.

            # the following types are added into dictionary:
            #   1. one-to-one aligned words (word -> word)
            #   2. one-to-many aligned words (word -> phrase)
            #
            # the following types are NOT added into dictionary:
            #   3. many-to-one aligned words (phrase -> word)
            #   4. many-to-many aligned words (phrase -> phrase)
    
    """
    
    en_fr = defaultdict(set)
    fr_en = defaultdict(set)
    for sent in aligned_sents:
        if len(sent.sent1_index) != len(sent.sent1) \
           or len(sent.sent2_index) != len(sent.sent2):
            raise IndexError("Alignment is incompletely defined.")
        for i in range(0, len(sent.sent1)):
            # add to dict iff the alignment is Type 1 or 2
            if [] == [x for x in \
                      [sent.sent2_index[y] for y in sent.sent1_index[i]]\
                      if len(x) != 1]:
                en_word = sent.sent1[i].lower()
                fr_words = ' '.join([ sent.sent2[x].lower() \
                                      for x in sorted(sent.sent1_index[i])])
                en_fr[en_word].add(fr_words)
            
        for i in range(0, len(sent.sent2)):
            # Ditto up
            if [] == [x for x in \
                      [sent.sent1_index[y] for y in sent.sent2_index[i]]\
                      if len(x) != 1]:
                fr_word = sent.sent2[i].lower()
                en_words = ' '.join([ sent.sent1[x].lower() \
                                      for x in sorted(sent.sent2_index[i])])
                fr_en[fr_word].add(en_words)
                
    return (en_fr, fr_en)

def main():
    """Test Case
    """
    # The simple example described in the bug tracker
    words = ['word1', 'word2', 'word3', 'word4']
    mots = ['mot1', 'mot2', 'mot3', 'mot4', 'mot5']
    alignments = [(0, 0), (1, 0), (2, 1), (3, 2), (3, 4)]
    aligned_sent = AlignedSent(words, mots, alignments)

    # Print words that align with forth word in words
    print [aligned_sent.sent2[n] for n in aligned_sent.sent1_index[3]]

    # Print out the GIZA form of this alignment
    print aligned_sent

    # Print word sets
    print aligned_sent.sent1_sets()
    print aligned_sent.sent2_sets()

    # Test gui_repr
    aligned_sent.gui_repr()
    AlignedSent.block_and_view_gui_repr()

if __name__ == "__main__":
    main()
